"""讯飞星火模型适配器实现。"""

from typing import AsyncGenerator, Dict, Any, List
import asyncio
import json
import hashlib
import hmac
import base64
from datetime import datetime
from urllib.parse import urlencode, urlparse
import websockets
import ssl

from .base import BaseModelAdapter, ModelConfig, ModelResponse


class SparkAdapter(BaseModelAdapter):
    """讯飞星火大模型适配器。"""
    
    # 定价（每千tokens，人民币）
    PRICING = {
        "spark-3.5": {"prompt": 0.002, "completion": 0.002},
        "spark-3.0": {"prompt": 0.002, "completion": 0.002},
        "spark-2.0": {"prompt": 0.001, "completion": 0.001}
    }
    
    # API版本映射
    API_VERSIONS = {
        "spark-3.5": "v3.5",
        "spark-3.0": "v3.1", 
        "spark-2.0": "v2.1"
    }
    
    def __init__(self, config: ModelConfig):
        """初始化讯飞适配器。
        
        需要在extra_params中提供:
        - app_id: 讯飞应用ID
        - api_secret: API密钥
        """
        super().__init__(config)
        self.app_id = config.extra_params.get("app_id", "")
        self.api_secret = config.extra_params.get("api_secret", "")
        self.api_version = self.API_VERSIONS.get(config.model_name, "v3.1")
        self.ws_url = self._get_ws_url()
    
    def _initialize_client(self) -> None:
        """初始化客户端（讯飞使用WebSocket）。"""
        # 讯飞使用WebSocket，无需初始化特定客户端
        pass
    
    def _get_ws_url(self) -> str:
        """生成WebSocket URL with HMAC-SHA256签名。"""
        # 讯飞星火WebSocket地址
        host = "spark-api.xf-yun.com"
        path = f"/{self.api_version}/chat"
        
        # 生成时间戳
        date = datetime.utcnow().strftime("%a, %d %b %Y %H:%M:%S GMT")
        
        # 构建签名原始字符串
        signature_origin = f"host: {host}\ndate: {date}\nGET {path} HTTP/1.1"
        
        # HMAC-SHA256签名
        signature = base64.b64encode(
            hmac.new(
                self.api_secret.encode(),
                signature_origin.encode(),
                hashlib.sha256
            ).digest()
        ).decode()
        
        # 构建authorization
        authorization_origin = f'api_key="{self.config.api_key}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature}"'
        authorization = base64.b64encode(authorization_origin.encode()).decode()
        
        # 构建请求参数
        params = {
            "authorization": authorization,
            "date": date,
            "host": host
        }
        
        # 返回完整URL
        return f"wss://{host}{path}?{urlencode(params)}"
    
    async def generate(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> ModelResponse:
        """通过WebSocket生成响应。
        
        Args:
            messages: 对话消息
            **kwargs: 额外参数
            
        Returns:
            模型响应
        """
        # 构建请求数据
        request_data = self._build_request(messages, **kwargs)
        
        # 连接WebSocket并获取响应
        full_response = ""
        usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
        
        async with websockets.connect(
            self.ws_url,
            ssl=ssl.create_default_context()
        ) as websocket:
            # 发送请求
            await websocket.send(json.dumps(request_data))
            
            # 接收响应
            while True:
                response = await websocket.recv()
                data = json.loads(response)
                
                if data.get("header", {}).get("code") != 0:
                    raise Exception(f"讯飞API错误: {data.get('header', {}).get('message', 'Unknown error')}")
                
                # 提取内容
                choices = data.get("payload", {}).get("choices", {})
                text = choices.get("text", [])
                if text:
                    content = text[0].get("content", "")
                    full_response += content
                
                # 提取使用统计
                usage_data = data.get("payload", {}).get("usage", {}).get("text", {})
                if usage_data:
                    usage["prompt_tokens"] = usage_data.get("prompt_tokens", 0)
                    usage["completion_tokens"] = usage_data.get("completion_tokens", 0)
                    usage["total_tokens"] = usage_data.get("total_tokens", 0)
                
                # 检查是否结束
                status = choices.get("status", 0)
                if status == 2:  # 2表示结束
                    break
        
        return ModelResponse(
            content=full_response,
            model=self.config.model_name,
            usage=usage,
            metadata={
                "app_id": self.app_id
            }
        )
    
    async def stream(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> AsyncGenerator[str, None]:
        """流式生成响应。
        
        Args:
            messages: 对话消息
            **kwargs: 额外参数
            
        Yields:
            内容块
        """
        request_data = self._build_request(messages, **kwargs)
        
        async with websockets.connect(
            self.ws_url,
            ssl=ssl.create_default_context()
        ) as websocket:
            # 发送请求
            await websocket.send(json.dumps(request_data))
            
            # 流式接收
            while True:
                response = await websocket.recv()
                data = json.loads(response)
                
                if data.get("header", {}).get("code") != 0:
                    raise Exception(f"讯飞API错误: {data.get('header', {}).get('message', 'Unknown error')}")
                
                # 提取并输出内容
                choices = data.get("payload", {}).get("choices", {})
                text = choices.get("text", [])
                if text:
                    content = text[0].get("content", "")
                    if content:
                        yield content
                
                # 检查是否结束
                status = choices.get("status", 0)
                if status == 2:
                    break
    
    def _build_request(self, messages: List[Dict[str, str]], **kwargs) -> Dict[str, Any]:
        """构建请求数据。
        
        Args:
            messages: 消息列表
            **kwargs: 额外参数
            
        Returns:
            请求数据字典
        """
        # 转换消息格式
        text_messages = []
        for msg in messages:
            role = msg["role"]
            if role == "system":
                role = "system"
            elif role == "assistant":
                role = "assistant"
            else:
                role = "user"
            
            text_messages.append({
                "role": role,
                "content": msg["content"]
            })
        
        return {
            "header": {
                "app_id": self.app_id,
                "uid": "user123"  # 可以使用实际用户ID
            },
            "parameter": {
                "chat": {
                    "domain": self.config.model_name.replace("-", "_"),
                    "temperature": self.config.temperature,
                    "max_tokens": self.config.max_tokens,
                    "top_k": kwargs.get("top_k", 4),
                    "chat_id": kwargs.get("chat_id", "default")
                }
            },
            "payload": {
                "message": {
                    "text": text_messages
                }
            }
        }
    
    def count_tokens(self, text: str) -> int:
        """计算文本的token数量。
        
        中文大约1个字符1个token。
        
        Args:
            text: 输入文本
            
        Returns:
            Token数量
        """
        # 中文字符计数
        chinese_chars = sum(1 for char in text if '\u4e00' <= char <= '\u9fff')
        # 英文单词计数
        english_words = len([w for w in text.split() if not any('\u4e00' <= c <= '\u9fff' for c in w)])
        
        return chinese_chars + int(english_words * 1.5)
    
    def estimate_cost(self, prompt_tokens: int, completion_tokens: int) -> float:
        """估算使用成本。
        
        Args:
            prompt_tokens: 提示词token数
            completion_tokens: 生成token数
            
        Returns:
            预估成本（美元）
        """
        model_key = self.config.model_name
        pricing = self.PRICING.get(model_key, self.PRICING["spark-3.0"])
        
        prompt_cost = (prompt_tokens / 1000) * pricing["prompt"]
        completion_cost = (completion_tokens / 1000) * pricing["completion"]
        
        # 转换为美元
        cny_to_usd = 0.14
        return (prompt_cost + completion_cost) * cny_to_usd